#include "classifier.h"
#include "fast-cpp-csv-parser/csv.h"
#include "boost/filesystem.hpp"
#include "caffe/util/db_lmdb.hpp"
#include "caffe/caffe.hpp"
#include "lmdb.h"
#include "tools.h"


using boost::filesystem::path;
using google::INFO;
using google::WARNING;
using google::ERROR;
using google::FATAL;
using cv::Scalar;
//using caffe::db::MDB_CHECK;
using caffe::db::GetDB;

inline void MDB_CHECK(int mdb_status) {
    CHECK_EQ(mdb_status, MDB_SUCCESS) <<
                                      mdb_strerror(mdb_status) <<
                                      Common::printTrace();
}
/*
 * Note that this macro set up the map environment;
 * Returns the reference of data in MAP;
 * Callers should user reference to change Data in MAP
 */
#ifndef CHECK_MAP_VAL
#define CHECK_MAP_VAL(key) \
    (\
        (\
            {\
                static string __registerName = registerFuncName(key); \
                if(imgStorage.find(key)!=imgStorage.end()){ \
                    return imgStorage[key];\
                }else{\
                    imgStorage[key] = CommonImage(__registerName);\
                }\
            }\
        ),\
        imgStorage[key] \
    )
#endif

/*
 * Consts used in this project
 */
int DRImage::NUM_DRIMAGE = 0;

ImgDataGenerator::ImgDataGenerator(string codeFilePath) {
    if (!boost::filesystem::exists(path(codeFilePath))) {
        LOG(WARNING) << "Cannot find file:" << codeFilePath << "Existing ...... ";
        return;
    }

    path filePath(codeFilePath);
    path parentPath = filePath.parent_path();

    io::CSVReader<3 + 1> in(codeFilePath);
    in.read_header(io::ignore_extra_column, "OriginalImage",
                   "Mas",
                   "Hard exclude",
                   "soft exclude"
    );
    string files[3 + 1];
    int id = 0;
    while (in.read_row(files[0], files[1], files[2], files[3])) {
        LOG(INFO) << "Reading : " << path(files[0]);
        ImageSet imageSet = ImageSet();
        // reading images from disk file
        for (int i = 0; i < Parameters::NUM_LABELS + 1; i++) {
            string imageFile = (parentPath / path(files[i])).string();
            if (!boost::filesystem::exists(path(imageFile))) {
                continue;
            }
            Mat tmp = cv::imread(imageFile);
            // resize images
            double ratio = (double) Parameters::IMAGE_COLS / tmp.cols;
            cv::resize(tmp, tmp, cv::Size(), ratio, ratio);
            if (i == 0) {
                imageSet.originalImage = DRImage(tmp, id);
            } else {
                Mat green = CommonImage(tmp).getGreenChannel().getOriginal();
                green = (green == 255);
                imageSet.labels.push_back(CommonImage(green, "label_" + Common::tostring(i)));
                imageSet.labels.back().write(id);
            }
        }
        // generating masks and interested pixels
        const Mat mask = imageSet.originalImage.getMask().getOriginal();
        imageSet.originalImage.getMask().write(id);
        imageSet.originalImage.write(id);
        for (int x = Parameters::WINDOW_SIZE / 2;
             x < mask.rows - Parameters::WINDOW_SIZE / 2;
             x += Parameters::NUM_STRIDE) {
            for (int y = Parameters::WINDOW_SIZE / 2;
                 y < mask.cols - Parameters::WINDOW_SIZE / 2;
                 y += Parameters::NUM_STRIDE) {
                if (mask.at<uchar>(x, y)) {
                    DataSample dataSample;
                    dataSample.imgIndex = id;
                    dataSample.row = x;
                    dataSample.col = y;
                    for (int numLabel = 0; numLabel < Parameters::NUM_LABELS; numLabel++) {
                        if (imageSet.labels[numLabel].getOriginal().at<uchar>(x, y)) {
                            dataSample.labels[numLabel] = 1;
                        } else {
                            dataSample.labels[numLabel] = 0;
                        }
                    }
                    dataSamples.push_back(dataSample);
                }
            }
        }
        images.push_back(imageSet);
        id++;
    }
}

Sample ImgDataGenerator::operator[](const int index) {
    int startRow = dataSamples[index].row - Parameters::WINDOW_SIZE / 2;
    int endRow = dataSamples[index].row + Parameters::WINDOW_SIZE / 2 + 1;
    int startCol = dataSamples[index].col - Parameters::WINDOW_SIZE / 2;
    int endCol = dataSamples[index].col + Parameters::WINDOW_SIZE / 2 + 1;
    int label = 0;
    for (int i = 0; i < Parameters::NUM_LABELS; i++)
        if (dataSamples[index].labels[i] > 0) {
            label = i + 1;
            break;
        }
    Mat pitch = images[dataSamples[index].imgIndex].
            originalImage.getOriginal().rowRange(startRow, endRow).
            colRange(startCol, endCol);
    return Sample{
            .data = pitch,
            .label = label
    };
}

size_t ImgDataGenerator::length() const {
    return dataSamples.size();
}


Classifier::Classifier(string trainFileList) :
        dg(trainFileList) {
}


string CommonImage::registerFuncName(string name) {
    if (registeredNames.find(name) != registeredNames.end()) {
        LOG(WARNING) << "Name : " << name << " already registered";
        LOG(WARNING) << Common::printTrace();
    }
    registeredNames[name] = 1;
    return name;
}

CommonImage::CommonImage(Mat img, string _name) :
        name(_name) {
    setImage(img);
}

CommonImage::CommonImage(string _name) :
        name(_name) { }

void CommonImage::setImage(Mat img) {
    CHECK(!img.empty()) << Common::printTrace();
    CHECK(img.depth() == CV_8U || img.depth() == CV_32F);
    CHECK(img.channels() == 1 || img.channels() == 3);
    imgStorage.clear();
    img.copyTo(originalImage);
}

CommonImage &CommonImage::getDouble() {
    auto &result = CHECK_MAP_VAL("double");
    CHECK(!getOriginal().empty());
    Mat tmp;
    if (getOriginal().depth() == CV_8U)
        getOriginal().convertTo(tmp, CV_32FC(getOriginal().channels()), 1.0 / 255.0);
    else if (getOriginal().depth() == CV_32F)
        tmp = getOriginal();
    else
        CHECK(false) << "ERRPR, Type of Image is: " << getOriginal().depth();
    result.setImage(tmp);
    return result;
}

CommonImage &CommonImage::getByte() {
    auto &result = CHECK_MAP_VAL("byte");
    CHECK(!getOriginal().empty());
    Mat tmp;
    if (getOriginal().depth() == CV_32F)
        getOriginal().convertTo(tmp, CV_8UC(getOriginal().channels()), 255.0);
    else if (getOriginal().depth() == CV_8U)
        tmp = getOriginal();
    else
        CHECK(false) << "ERRPR, Type of Image is: " << getOriginal().depth();
    result.setImage(tmp);
    return result;
}

CommonImage &CommonImage::getChannel() {
    auto &G = CHECK_MAP_VAL("green_channel");
    auto &R = CHECK_MAP_VAL("red_channel");
    auto &B = CHECK_MAP_VAL("blue_channel");
    vector<Mat> channels;
    cv::split(originalImage, channels);
    CHECK_EQ(3, channels.size());
    R.setImage(channels[0]);
    G.setImage(channels[1]);
    B.setImage(channels[2]);
    return G;
}

CommonImage &CommonImage::getGreenChannel() {
    getChannel();
    if (imgStorage.find("green_channel") != imgStorage.end()) {
        return imgStorage["green_channel"];
    }
    CHECK(false) << "Green Channel should already exist!";
    LOG(INFO) << "FUCK";
    return imgStorage["green_channel"];
}


const Mat &CommonImage::getOriginal() const {
    CHECK(!originalImage.empty());
    return originalImage;
}

void CommonImage::show(string name, int pause) const {
    static int counter = 1;
    if (name == "")
        name = string("image_") + Common::tostring(counter++);
    name = this->name + "_" + name;
    cv::namedWindow(name.c_str(), CV_WINDOW_NORMAL);
    cv::imshow(name.c_str(), getOriginal());
    cv::resizeWindow(name, 800, 600);
    cv::waitKey(pause);
}

/*
 * Write image to disk files.
 * If no name and id given, than use default name in class;
 */
void CommonImage::write(int id) const {
    string prefix = "./img/";
    path fpath = path(prefix);
    if (!boost::filesystem::exists(fpath)) {
        boost::filesystem::create_directories(fpath);
    }
    string _name = name;
    if (_name == "") {
        _name = "Figure_";
    }
    if (id >= 0) {
        _name = "Figure_" + Common::tostring(id) + "_" + _name + ".png";
    } else {
        _name = _name + ".png";
    }
    fpath = fpath / path(_name);
    cv::imwrite(fpath.string().c_str(), originalImage);
}


DRImage::DRImage(Mat img, int _id) : CommonImage("OriginalImg"), id(--NUM_DRIMAGE) {
    // perform preprocess
//    vector<Mat> channels;
//    cv::split(img, channels);
//    channels[0] = CommonImage(channels[1]).getDouble().getLocalNorm(20, 80).getByte().getOriginal();
    Mat newImg = img;
//    cv::cvtColor(img,newImg,cv::COLOR_RGB2HSV);
//    cv::merge(channels, newImg);
    setImage(newImg);
    if (_id >= 0)
        id = _id;
}

DRImage::DRImage(int _id) : CommonImage("OriginalImg"), id(--NUM_DRIMAGE) {
    if (_id >= 0)
        id = _id;
}

/*
 *  Mask CV8UC1
 */
const CommonImage &DRImage::getMask(double thresh) {
    auto &result = CHECK_MAP_VAL("mask");
    Mat img = getDouble().getOriginal();
    vector<Mat> image_channels;
    split(img, image_channels);
    Mat rChannel(img.rows, img.cols, CV_32FC1, Scalar(0));
    for (int i = 1; i < 2; i++) {
        rChannel += image_channels[i].mul(image_channels[i]);
    }
    double maxValue, minValue;
    minMaxLoc(rChannel, &minValue, &maxValue);
    result.setImage(rChannel > thresh * maxValue);
    CHECK(!result.getOriginal().empty());
    CHECK_EQ(CV_8U, result.getOriginal().depth());
    CHECK_EQ(1, result.getOriginal().channels());
    return result;
}

const CommonImage &DRImage::getHardLeakageNoneCaffe() {
    auto &result = CHECK_MAP_VAL("hardLeakage");

    return result;
}

const CommonImage &DRImage::getIgnoreFiel() {
    auto &result = CHECK_MAP_VAL("ignoreField");

    return result;
}


LMDBDataSet::LMDBDataSet() : size(0), keySize(0) {
}

/*
 * This function is another entrance of this class.
 * This function set up a write environment.
 */
void LMDBDataSet::newDB(string DBPath) {
    MDB_CHECK(mdb_env_create(&mdb_env_));
    CHECK(
            boost::filesystem::create_directories(path(DBPath))
    ) << Common::printTrace() << "make directory " << DBPath << "failed";
    unsigned int flags = 0;
    MDB_CHECK(mdb_env_open(mdb_env_, DBPath.c_str(), flags, 0664));
    LOG(INFO) << "Opened lmdb " << DBPath;

    transaction = newWriteTransaction(mdb_env_);
    keySize = 10;
    currentPosition = 0;
}

size_t LMDBDataSet::getKeySize() {
    MDB_txn *mdb_txn_;
    MDB_dbi mdb_dbi_;
    MDB_cursor *mdb_cursor;
    MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
    MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
    MDB_CHECK(mdb_cursor_open(mdb_txn_, mdb_dbi_, &mdb_cursor));
    MDB_val key, val;
    key.mv_size = 0;
    key.mv_data = NULL;
    MDB_CHECK(mdb_cursor_get(mdb_cursor, &key, &val, MDB_FIRST));
    mdb_cursor_close(mdb_cursor);
    mdb_dbi_close(mdb_env_, mdb_dbi_);
    mdb_txn_abort(mdb_txn_);
    return (int) key.mv_size;
}


/*
 * This function is the actual a entrance of this class,
 * It initialise the read environment.
 */
bool LMDBDataSet::loadFromLMDB(string DBPath) {
    MDB_CHECK(mdb_env_create(&mdb_env_));
    unsigned int flags = 0;
    flags = MDB_RDONLY | MDB_NOTLS;
    int rc = mdb_env_open(mdb_env_, DBPath.c_str(), flags, 0664);
    MDB_CHECK(rc);
    LOG(INFO) << "Opened lmdb " << DBPath;
    keySize = getKeySize();
    LOG(INFO) << "got key size " << keySize;
    // set up transactions
    transaction = newReadTransaction(mdb_env_);
    return true;
}

size_t LMDBDataSet::length() {
    if (size)
        return size;
    CHECK_GT(keySize, 0) <<
                         Common::printTrace() <<
                         "Key size Not Set, get key size First";
    size_t maxsize = UINT_MAX / 2;
    size_t minsize = 0;
    MDB_txn *mdb_txn_;
    MDB_dbi mdb_dbi_;
    MDB_val key, val;
    MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, MDB_RDONLY, &mdb_txn_));
    MDB_CHECK(mdb_dbi_open(mdb_txn_, NULL, 0, &mdb_dbi_));
    while (minsize < maxsize) {
        size_t half = (minsize + maxsize) / 2;
        string tmp = caffe::format_int((int) half, (int) keySize);
        key.mv_size = tmp.size();
        key.mv_data = (void *) tmp.data();
        int res = mdb_get(mdb_txn_, mdb_dbi_, &key, &val);
        if (res == 0) {
            minsize = half + 1;
        } else if (res == MDB_NOTFOUND) {
            maxsize = half;
        } else {
            MDB_CHECK(res);
        }
    }
    if (mdb_txn_ != NULL)
        mdb_txn_abort(mdb_txn_);
    if (mdb_dbi_)
        mdb_dbi_close(mdb_env_, mdb_dbi_);
    size = minsize;
    return size;
}

LMDBDataSet::~LMDBDataSet() {
    close();
}

bool LMDBDataSet::seek(int index) {
    currentPosition = index;
    valid_ = currentPosition < length();
    return valid_;
}

bool LMDBDataSet::getDatum(caffe::Datum &datum) {
    return datum.ParseFromString(getMDB());
}

string  LMDBDataSet::getMDB() {
    return
            transaction->get(caffe::format_int(currentPosition, (int) keySize));
}

Mat LMDBDataSet::datumToMat(caffe::Datum *datum) {
    Mat mat;
    int datum_channels = datum->channels();
    int datum_height = datum->height();
    int datum_width = datum->width();
    int datum_size = datum_channels * datum_height * datum_width;

    if (datum->data().size() > 0) {
        CHECK_EQ(datum->data().size(), datum_size);
        const uchar *const dataFromDatum = (const uchar *const) (datum->data().c_str());
        mat = Mat(datum->width(), datum->height(), CV_8UC((datum->channels())));
        for (int h = 0; h < datum_height; h++) {
            uchar *ptr = mat.ptr(h);
            int img_index = 0;
            for (int w = 0; w < datum_width; w++) {
                for (int c = 0; c < datum_channels; c++) {
                    int datum_index = (c * datum_height + h) * datum_width + w;
                    ptr[img_index++] = dataFromDatum[datum_index];
                }
            }
        }
    } else if (datum->float_data_size() > 0) {
        LOG(INFO) << "getting float data;";
        LOG(FATAL) << "Converting float datum is not implemented yet";
    }
    return mat;
}

Mat LMDBDataSet::getCvMat() {
    caffe::Datum datum;
    getDatum(datum);
    return datumToMat(&datum);
}

Sample LMDBDataSet::getSample() {
    caffe::Datum datum;
    CHECK(getDatum(datum));
    Sample sample{
            .data = datumToMat(&datum),
            .label = datum.label()
    };
    return sample;
}

/*
 * This method is implemented using the cursor interface;
 * So this functionn will interfere with cursor.
 * Don't use this function with next() function or other
 * cursor interface
 */
Sample LMDBDataSet::operator[](int index) {
    CHECK(index < length()) <<
                            "Length is " << length() <<
                            " But Asking " << index <<
                            Common::printTrace();
    seek(index);
    return getSample();
}

/*
 * This function is designed to get user defined index or use
 * self increment index.
 * Not that only one form of index method should be used.
 * Otherwise Duplicated key may occur in lmdb.
 *
 * Also, User should avoid duplicate key when using user
 * defined keys.
 */
void LMDBDataSet::addSample(Sample sample, int index) {
    caffe::Datum datum;
    // TODO change this to my own code and add float type convert
    caffe::CVMatToDatum(sample.data, &datum);
    datum.set_label(sample.label);
    if (index == -1) {
        index = currentPosition;
        currentPosition++;
    }
    while (existedKeys.find(index) != existedKeys.end()) {
        LOG(WARNING) << "Key " << index << " is already in use";
        index++;
    }
    existedKeys[index] = true;
    addDatum(index, &datum);
}

void LMDBDataSet::addDatum(int key, caffe::Datum *datum) {
    string data;
    CHECK(datum->SerializeToString(&data));
    string key_ = caffe::format_int(key, (int) keySize);
    put(key_, data);
}

void LMDBDataSet::close() {
    if (transaction) {
        transaction->commit();
        delete transaction;
        transaction = NULL;
    }
    if (mdb_env_)
        mdb_env_close(mdb_env_);
    mdb_env_ = NULL;
    size = 0;
    keySize = 0;
}

CommonImage &CommonImage::getLocalNorm(float sigma1, float sigma2, const cv::_InputArray &mask) {
    auto &result = CHECK_MAP_VAL("localNorm");
    CHECK_EQ(this->originalImage.channels(), 1);
    Mat image = this->getDouble().getOriginal();
    Mat floatGray, blurred1, blurred2, temp1, temp2, res;
//    image.convertTo(floatGray, CV_32FC1);
    CHECK(image.channels() == 1);
    CHECK(image.depth() == CV_32F);
    image.copyTo(floatGray);

    int blur1 = (int) (2 * ceil(-NormInv(0.05, 0, sigma1)) + 1);
    cv::GaussianBlur(floatGray, blurred1, cv::Size(blur1, blur1), sigma1);
    temp1 = floatGray - blurred1;

    cv::pow(temp1, 2.0, temp2);
    int blur2 = (int) (2 * ceil(-NormInv(0.05, 0, sigma2)) + 1);
    cv::GaussianBlur(temp2, blurred2, cv::Size(blur2, blur2), sigma2);
    cv::pow(blurred2, 0.5, temp2);

    floatGray = temp1 / temp2;
//    floatGray = 255.0*floatGray;
//    floatGray.convertTo(res, CV_8UC1);
    normalize(floatGray, res, 0, 1, cv::NORM_MINMAX, -1, mask);
    result.setImage(res);
    return result;
}


